from typing import Annotated,Literal

import os

from autogen import ConversableAgent

Operator = Literal["+","-","*","/"]

# 配置中文LLM
llm_config = {
    "config_list": [
        {
            "model": "qwen-turbo",
            "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
            "api_key":os.environ["DASHSCOPE_API_KEY"]
        },
        # {"model": "gpt-4", "api_key": os.environ["OPENAI_API_KEY"]}
    ],
}

def calculator(a:int,b:int,operator:Annotated[Operator,"operator"]) -> int:
    if operator == "+":
        return a + b
    elif operator == "-":
        return a - b
    elif operator == "*":
        return a * b
    elif operator == "/":
        return int(a / b)
    else:
        raise ValueError("Invalid operator")


# Let's first define the assistant agent that suggests tool calls.
assistant = ConversableAgent(
    name="Assistant",
    system_message="You are a helpful AI assistant. "
    "You can help with simple calculations. "
    "Return 'TERMINATE' when the task is done.",
    llm_config=llm_config
)

# The user proxy agent is used for interacting with the assistant agent
# and executes tool calls.
user_proxy = ConversableAgent(
    name="User",
    llm_config=False,
    is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
    human_input_mode="NEVER",
)

# Register the tool signature with the assistant agent.
assistant.register_for_llm(name="calculator", description="A simple calculator")(calculator)

# Register the tool function with the user proxy agent.
user_proxy.register_for_execution(name="calculator")(calculator)

# 也可以使用下面的方法一次向两个代理注册工具
# from autogen import register_function
# # Register the calculator function to the two agents.
# register_function(
#     calculator,
#     caller=assistant,  # The assistant agent can suggest calls to the calculator.
#     executor=user_proxy,  # The user proxy agent can execute the calculator calls.
#     name="calculator",  # By default, the function name is used as the tool name.
#     description="A simple calculator",  # A description of the tool.
# )

chat_result = user_proxy.initiate_chat(assistant, message="(44232+13312/(232-32))*5 的结果是什么？")

